#!/usr/bin/python3
# -*- coding:utf-8 -*-
import argparse
import multiprocessing
import os
import sys

sys.path.insert(90, os.path.dirname(__file__) + "/../lib/redis2")
from lib import comred
from lib import interactive

if sys.version.startswith("2"):
    from lib.log import log2 as log
else:
    from lib.log import log

global localip

def initLocal(lhost):
    global localip
    localip = lhost

def getargs():
    parser = argparse.ArgumentParser(
        description='get redis shell.')  # 获取解析器对象,description为使用-h参数时输出的描述信息
    parser.add_argument('-r', '--rhost', dest='rhost', type=str, help='a single target', action='store')
    parser.add_argument('-p', '--rport', dest='rport', type=int, help='the port of remote redis server.')
    parser.add_argument('-f', '--file', dest='file', type=str, help='attach all targets from file.', action='store')
    parser.add_argument("-a", "--auth", dest="auth", type=str, help="password.", action='store')
    parser.add_argument('-l', '--lhost', dest='lhost', type=str, help='localhost ip.')
    return parser


def c(lhost, rhost, rport=6379, rshells={}, auth=None):
    conredis = comred.conredis(lhost)
    if rport == None:
        rport = 6379
    if auth != None:
        if conredis.connect(rhost, rport, auth) == None:
            log.mlog.prin("[!] Redis connect failed:" + rhost + ":" + str(rport))
            return False

    elif conredis.connect(rhost, rport) == None:
        log.mlog.prin("[!] Redis connect failed:" + rhost + ":" + str(rport))
        return False
    rshells[rhost] = conredis


def run():
    parser = getargs().parse_args()
    shells = multiprocessing.Manager().dict()
    rshells = {}

    if parser.lhost == None:
        log.mlog.prin("-l/--lhost cant be empty.")
        exit(0)
    initLocal(parser.lhost)
    if parser.rhost != None and parser.file == None:
        if c(parser.lhost, parser.rhost, parser.rport, rshells, parser.auth) == False:
            exit(1)
        interactive.interactive(shells, rshells)
        exit(0)
    elif parser.file != None and parser.rhost == None:
        with open(parser.file, "r") as f:
            rhosts = f.readlines()
        for r in rhosts:
            tmp = []
            slp = " "
            if "\t" in r:
                slp = "\t"
            for t in r.strip().split(slp):
                if t != "" and t != " ":
                    tmp.append(t)

            if len(tmp) == 1:
                ab = c(parser.lhost, r.strip(), 6379, rshells, parser.auth)
            if len(tmp) == 2:
                ab = c(parser.lhost, tmp[0], int(tmp[1]), rshells, parser.auth)
            if len(tmp) == 3:
                ab = c(parser.lhost, tmp[0], int(tmp[1]), rshells, tmp[2])
            if ab == False:
                continue

        interactive.interactive(shells, rshells)
        exit(0)
    else:
        log.mlog.prin("Only one of -r/-f can and must be used at the same time.")
        exit(0)
